#include <time.h>
#include <chrono>
#include <cmath>
#include <fstream>
#include <iostream>
#include <random>

#include "glog/logging.h"

// gtsam
#include "gtsam/base/Matrix.h"
#include "gtsam/base/Vector.h"

#include "gtsam/inference/Key.h"
#include "gtsam/inference/Symbol.h"
#include "gtsam/nonlinear/ISAM2.h"
#include "gtsam/nonlinear/NonlinearFactorGraph.h"
#include "gtsam/nonlinear/Values.h"

// expression: auto diff
#include "gtsam/nonlinear/ExpressionFactorGraph.h"
#include "gtsam/nonlinear/expressions.h"

// function: y = ax*x + bx + c
gtsam::Vector1 CurveFitting(const gtsam::Vector3& params,
                            const double x,
                            gtsam::OptionalJacobian<1, 3> H = boost::none,
                            gtsam::OptionalJacobian<1, 1> H2 = boost::none) {
  // jacobian
  if (H) {
    (*H) = (gtsam::Matrix(1, 3) << x * x, x, 1.0).finished();
  }
  double y_pred = params[0] * x * x - params[1] * x + params[2];
  return (gtsam::Vector(1) << y_pred).finished();
}

typedef gtsam::Expression<gtsam::Vector1> Vector1_;
typedef gtsam::Expression<gtsam::Vector3> Vector3_;
typedef gtsam::Expression<double> double_;

// expresion function: y = ax*x + bx + c
Vector1_ CurveFitting_(const Vector3_& params, const double_ x) {
  return Vector1_(&CurveFitting, params, x);
}

int main(int argc, char** argv) {
  // random gen
  uint32_t curr_time = static_cast<uint32_t>(std::time(nullptr));
  std::random_device rd{};
  std::mt19937 gen{rd()};
  LOG(INFO) << "random generator seed: " << curr_time;
  gen.seed(curr_time);
  std::normal_distribution<double> normal_dist{0.0, 5};
  // generate ground truth dataset
  std::vector<double> x_datas, y_datas;
  double params[] = {0.00332434324, 0.2133333, 4.311111};

  std::ofstream data_fout("/tmp/sim_data.txt");
  data_fout.precision(10);
  for (int i = -500; i < 500; i++) {
    x_datas.push_back(i);
    double y_noise =
        params[0] * i * i + params[1] * i + params[2] + normal_dist(gen);
    y_datas.push_back(y_noise);
    data_fout << i << " " << y_noise << std::endl;
  }
  data_fout.close();

  // isam
  gtsam::ISAM2Params isam2_params;
  isam2_params.relinearizeSkip = 100;
  isam2_params.relinearizeThreshold = 0.1;
  gtsam::ISAM2 isam2_solver(isam2_params);

  gtsam::ExpressionFactorGraph graph;
  gtsam::Values init_estimates;
  gtsam::Values results;

  // data noise model
  auto data_noise_model =
      gtsam::noiseModel::Diagonal::Sigmas(gtsam::Vector1(5));

  Vector3_ params_(gtsam::Symbol('x', 0));

  // build factor graph
  for (int i = 0; i < y_datas.size(); i++) {
    graph.addExpressionFactor(CurveFitting_(params_, x_datas[i]),
                              gtsam::Vector1(y_datas[i]),
                              data_noise_model);
  }
  // init vaules for params
  init_estimates.insert<gtsam::Vector3>(gtsam::Symbol('x', 0),
                                        gtsam::Vector3(0.0, 0.0, 0.0));
  // optimization
  auto opt_start = std::chrono::steady_clock::now();
  isam2_solver.update(graph, init_estimates);
  results = isam2_solver.calculateEstimate();
  auto opt_end = std::chrono::steady_clock::now();
  double opt_dt = std::chrono::duration_cast<std::chrono::duration<double>>(
                      opt_end - opt_start)
                      .count();
  LOG(INFO) << "iSAM2 optimize duration: " << opt_dt;
  LOG(INFO) << "results: "
            << results.at<gtsam::Vector3>(gtsam::Symbol('x', 0)).transpose();
}
